# _*_ coding:utf-8 _*_
# 开发团队：huang
# 开发时间：2020/11/19 22:50
# 文件名称：function.py
# 开发工具：PyCharm

"""
    1. video_insert_watermark：向视频中插入数字水印
    2. video_add_visual_watermark：对视频进行压缩再插入明水印
    3. 找不到probe的话需要先安装 ffmpeg
"""
from method import *


def video_insert_watermark( input_filepath, watermark_filepath, output_filepath ):
    """
    1. 创建一个 video_name_tmp 文件夹用于存储相应的中间文件，最后一步需要删除这个文件夹
    2. 将视频转换为 mp4格式，再将视频流和音频流分开保存
    3. 将视频流文件分成多个 5 分钟的片段
    4. 往每个片段前50帧添加水印
    :param output_filepath:
    :param input_filepath:
    :param watermark_filepath:
    :return:
    """

    # 1. 创建一个临时文件夹用来存储相应的临时文件
    # tmp_directory = '/home/data/video/target/filename_tmp'
    (_, tmp_filename) = os.path.split(input_filepath)  # 分离路径和文件名
    tmp_name, _ = os.path.splitext(tmp_filename)  # 分离文件名和后缀
    (tmp_path, _) = os.path.split(output_filepath)
    tmp_directory = "{0}/{1}_tmp".format(tmp_path, tmp_name)
    create_dir(tmp_directory)

    # 2. 将视频转换为 mp4格式
    # 转换后的视频文件路径为 reformat_filepath = '/home/data/video/target/filename_tmp/reformat.mp4'
    reformat_filepath = "{0}/reformat.mp4".format(tmp_directory)
    reformat2mp4(input_filepath, reformat_filepath)

    # 取出视频流的多个参数不用每次都取提高速度
    video_param = get_video_param(reformat_filepath)

    # 判断是否存在音频流，存在的话将其分离出来存储到临时文件夹中
    if video_param['audio']:
        audio_stream_path = "{0}/audio.m4a".format(tmp_directory)
        os.system("ffmpeg -y -i {0} -acodec copy -vn {1}".format(reformat_filepath, audio_stream_path))

    # 视频流存储路径
    video_stream_path = "{0}/video.mp4".format(tmp_directory)
    os.system("ffmpeg -y -i {0} -vcodec copy -an {1}".format(reformat_filepath, video_stream_path))

    delete_file(reformat_filepath)

    # 将视频切分成多个 5 分钟的片段，每一个片段开一个进程进行操作
    video_stream_embed_path = "{0}/video_embed.mp4".format(tmp_directory)
    split_video(video_stream_path, video_stream_embed_path, watermark_filepath, video_param)

    delete_file(video_stream_path)

    # 音频视频合并
    tmp_ouput_filepath = "{0}/tmp_output.mp4".format(tmp_directory)
    if video_param['audio']:
        os.system("ffmpeg -y -i {0} -i {1} -c copy {2}".format(video_stream_embed_path, audio_stream_path,
                                                               tmp_ouput_filepath))
    else:
        os.system("cp {0} {1}".format(video_stream_embed_path, tmp_ouput_filepath))

    delete_file(video_stream_embed_path)

    # 将生成的视频复制到 output 中，增加这一步可以避免未完全生成完合并文件就使用输出文件
    os.system("cp {0} {1}".format(tmp_ouput_filepath, output_filepath))
    delete_file(tmp_ouput_filepath)

    # 清空临时文件夹
    delete_dir(tmp_directory)


def video_extract_watermark( input_filepath, output_filepath ):
    """
    找一个超过 10s 的片段，判断是否存在水印
    :param input_filepath:
    :param output_filepath:
    :return:
    """
    # 创建存储中间文件的文件夹
    # 提取出的 watermark 所在文件夹
    (_, tmp_filename) = os.path.split(input_filepath)  # 分离路径和文件名
    tmp_name, _ = os.path.splitext(tmp_filename)  # 分离文件名和后缀
    (tmp_path, _) = os.path.split(output_filepath)
    tmp_directory = "{0}/{1}_tmp".format(tmp_path, tmp_name)
    create_dir(tmp_directory)

    # 1. 转为 mp4 格式

    # 2. 提取出视频流
    video_stream_path = "{0}/video.mp4".format(tmp_directory)
    os.system("ffmpeg -y -i {0} -vcodec copy -an {1}".format(input_filepath, video_stream_path))

    # 3. 截取前 20s 视频
    intra_filepath = "{0}/intra.mp4".format(tmp_directory)
    os.system("ffmpeg -y -i %s -strict -2 -qscale 0 -intra %s" % (video_stream_path, intra_filepath))

    gap_path = "{0}/gap.mp4".format(tmp_directory)
    os.system(
        "ffmpeg -y -ss {0} -t {1} -i {2} -vcodec copy -an {3}".format(0, 20, intra_filepath, gap_path))

    # 转为 yuv 格式
    video_param = get_video_param(gap_path)
    gap_yuv_path = "{0}/{1}.yuv".format(tmp_directory, tmp_name)
    os.system("ffmpeg -y -vcodec {0} -i {1} -pix_fmt {2} {3}".format(video_param['codec_name'], gap_path,
                                                                     video_param['pix_fmt'], gap_yuv_path))

    # 查找 yuv 文件中是否存在 flag
    exist_watermark = find_flag(gap_yuv_path, (video_param['height'], video_param['width']), output_filepath)
    delete_dir(tmp_directory)

    return exist_watermark


def video_add_visual_watermark( input_filepath, watermark_path, output_filepath ):
    # 1. 对视频进行压缩
    (_, tmp_filename) = os.path.split(input_filepath)  # 分离路径和文件名
    (filepath, _) = os.path.split(output_filepath)

    compress_output_filepath = "{0}/compress_{1}".format(filepath, tmp_filename)  # 压缩视频路径
    only_center_filepath = "{0}/only_center_{1}".format(filepath, tmp_filename)  # 只插入一个水印的路径
    two_watermark_filepath = "{0}/two_watermark_{1}".format(filepath, tmp_filename)  # 插入两个水印的路径

    compress_video(input_filepath, compress_output_filepath)  # 压缩视频

    # 读取视频的长宽
    video_param = get_video_param(compress_output_filepath)
    width = video_param['width']
    height = video_param['height']

    # 读取水印的大小，水印大小固定
    watermark = cv2.imread(watermark_path)
    watermark_width = watermark.shape[1]
    watermark_height = watermark.shape[0]

    # 2. 在中间加入水印 的坐标
    pos_x = int((width - watermark_width) / 2)
    pos_y = int((height - watermark_height) / 2)

    os.system(
        "ffmpeg -y -i {0} -i {1} -filter_complex \"[1:v]colorkey=0x000000:0.5:1.0[logo];[0:v][logo]overlay=x={2}:y={3}\" {4}".format(
            compress_output_filepath,
            watermark_path,
            pos_x, pos_y,
            only_center_filepath))

    # 删除使用完不再需要的中间文件
    delete_file(compress_output_filepath)

    # 判断插入第二个水印的坐标
    if width <= height:
        # 若高大于宽则在左上角插入水印
        pos_x = 5
        pos_y = 5

    else:
        # 若宽大于高则在右下角插入水印
        pos_x = width - watermark.shape[1] - 5
        pos_y = height - watermark.shape[0] - 5

    # 插入第二个水印
    os.system(
        "ffmpeg -y -i {0} -i {1} -filter_complex \"[1:v]colorkey=0x000000:0.5:1.0[logo];[0:v][logo]overlay=x={2}:y={3}\" {4}".format(
            only_center_filepath,
            watermark_path,
            pos_x, pos_y,
            two_watermark_filepath))

    # 删除使用完不再需要的中间文件
    delete_file(only_center_filepath)

    # 将生成的文件复制到目标文件中，防止插入水印时间过长目标文件使用出错
    os.system("cp {0} {1}".format(two_watermark_filepath, output_filepath))
    delete_file(two_watermark_filepath)
